#!/usr/bin/env python3

import rospy
import sys
import rosbag
from message_filters import ApproximateTimeSynchronizer, Subscriber
from sensor_msgs.msg import PointCloud2
from sensor_msgs import point_cloud2
import open3d as o3d
from nav_msgs.msg import Odometry
import os

# 默认参数值
rosfilepath = "/workspace/bags/hagar_mid360.bag"
outputfolder = "/workspace/bags/hagar_mid360/"

def callback_sync(lidar_msg, odom_msg):
    # print("Received synchronized messages!")
    lidar_timestamp = lidar_msg.header.stamp
    # print(f"odom_timestamp time: {lidar_timestamp}")
    odom_timestamp = odom_msg.header.stamp
    # print(f"odom_timestamp time: {odom_timestamp}")
    
    # Convert PointCloud2 message to Open3D PointCloud
    pc_data = point_cloud2.read_points(lidar_msg)
    points = []
    for p in pc_data:
        points.append([p[0], p[1], p[2]])
        
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)

    # Save PointCloud as PCD file
    o3d.io.write_point_cloud(f'{outputfolder}/{lidar_timestamp}.pcd', pcd)
    
    pose_position = odom_msg.pose.pose.position
    orientation = odom_msg.pose.pose.orientation

    # 将信息写入文本文件
    with open(f'{outputfolder}/{lidar_timestamp}.pos', 'w') as file:
        file.write(f"Position: {pose_position.x}, {pose_position.y}, {pose_position.z}\n")
        file.write(f"Orientation: {orientation.x}, {orientation.y}, {orientation.z}, {orientation.w}\n")
 
def writePoints(lidar_msg):
    lidar_timestamp = lidar_msg.header.stamp
    # Convert PointCloud2 message to Open3D PointCloud
    pc_data = point_cloud2.read_points(lidar_msg)
    points = []
    for p in pc_data:
        points.append([p[0], p[1], p[2]])
        
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    # Save PointCloud as PLY file
    o3d.io.write_point_cloud(f'{outputfolder}/{lidar_timestamp}_global.ply', pcd)
    
def main():
    # 初始化ROS节点
    rospy.init_node('transformbag2pcd', anonymous=True)

    # 获取命令行参数
    args = rospy.myargv(argv=sys.argv)

    global rosfilepath
    global outputfolder
    
    # 解析命令行参数
    if len(args) > 2:
        rosfilepath = args[1]
        outputfolder = args[2]

    # 输出参数值
    rospy.loginfo("rosfilepath %s", rosfilepath)
    rospy.loginfo("outputfolder %s", outputfolder)
    os.makedirs(outputfolder, exist_ok=True)

    
    # read bag
    with rosbag.Bag(rosfilepath, 'r') as bag:
        # 打印bag文件信息
        print(f"ROS Bag File: {rosfilepath}")
        print(f"Start Time: {bag.get_start_time()}")
        print(f"End Time: {bag.get_end_time()}")
        print(f"Topics: {bag.get_type_and_topic_info().topics}")
        
        # 创建消息同步器
        lidar_sub = Subscriber('/quad0_pcl_render_node/sensor_cloud', PointCloud2)
        odom_sub = Subscriber('/quad_0/lidar_slam/odom', Odometry)
        sync = ApproximateTimeSynchronizer([lidar_sub, odom_sub], queue_size=100, slop=0.01)
        sync.registerCallback(callback_sync)

        for topic, msg, t in bag.read_messages():
            if topic == '/quad0_pcl_render_node/sensor_cloud':
                lidar_sub.signalMessage(msg)
            elif topic == '/quad_0/lidar_slam/odom':
                odom_sub.signalMessage(msg)
            # elif topic == '/quad0_pcl_render_node/cloud':
            #     writePoints(msg)

if __name__ == '__main__':
    try:
        main()
    except rospy.ROSInterruptException:
        pass